// Copyright 2025 The XLS Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef XLS_CODEGEN_VAST_DSLX_TYPE_FIXER_H_
#define XLS_CODEGEN_VAST_DSLX_TYPE_FIXER_H_

#include <memory>

#include "xls/dslx/frontend/ast_cloner.h"
#include "xls/dslx/frontend/module.h"
#include "xls/dslx/import_data.h"
#include "xls/dslx/type_system/type_info.h"
#include "xls/dslx/type_system_v2/type_inference_error_handler.h"

namespace xls {

// A utility used internally by a `DslxBuilder` to perform cleanup on a draft
// AST generated by VAST->DSLX translation.
class DslxTypeFixer {
 public:
  virtual ~DslxTypeFixer() = default;

  // Returns an error handler that can fix type errors in the draft of the
  // generated DSLX. This must be used in 2 steps:
  // 1. Run type inference, passing this handler to ParseAndTypecheck.
  // 2. Clone the AST using GetErrorFixReplacer() to inject any needed fixes
  //    that were discovered during the type inference (the error handler itself
  //    does not change the AST during type inference).
  virtual dslx::TypeInferenceErrorHandler GetErrorHandler() = 0;

  // Deals out a `CloneReplacer` that can be used to transform the draft AST
  // into a fixed AST (i.e. inject needed casts and annotations).
  virtual dslx::CloneReplacer GetErrorFixReplacer(const dslx::TypeInfo* ti) = 0;

  // Deals out a `CloneReplacer` that simplifies a draft of the AST which may
  // have overuse of casts and explicit annotations. Typically the first draft
  // of the AST that we produce has aggressive/pessimistic casting, and the
  // error fix run adds more that it missed. A pass with the simplify replacer
  // should be run after the error fix pass so its input will be the AST with
  // "maximal casting".
  virtual dslx::CloneReplacer GetSimplifyReplacer(const dslx::TypeInfo* ti) = 0;
};

// Deals out a type fixer.
std::unique_ptr<DslxTypeFixer> CreateDslxTypeFixer(
    dslx::Module& original_module, const dslx::ImportData& import_data);

}  // namespace xls

#endif  // XLS_CODEGEN_VAST_DSLX_TYPE_FIXER_H_
